import pandas as pd
from glob import glob
import os
import warnings
from joblib import Parallel, delayed

pd.set_option('display.max_rows', 1000)
pd.set_option('expand_frame_repr', False)  # 当列太多时不换行
warnings.filterwarnings("ignore")

if_use_spot = True
n_jobs = 59

ROOT_PATH = os.path.abspath(os.path.join(os.path.abspath(os.path.dirname(__file__)) , '..'))  # 返回根目录文件夹
xbx_data_path = os.path.join(ROOT_PATH, 'data', 'xbx_data')
pickle_path = os.path.join(ROOT_PATH, 'data', 'pickle_data')
swap_path = os.path.join(pickle_path, 'swap')
spot_path = os.path.join(pickle_path, 'spot')
if not os.path.exists(pickle_path):
    os.makedirs(pickle_path)
if not os.path.exists(swap_path):
    os.makedirs(swap_path)
if not os.path.exists(spot_path):
    os.makedirs(spot_path)

stable_symbol = ['BKRW', 'USDC', 'USDP', 'TUSD', 'BUSD', 'FDUSD', 'DAI', 'EUR', 'GBP', 'USBP', 'SUSD', 'PAXG', 'AEUR']

# 若分析数据完整性有新的symbol报错日内数据不完整，需要更新SETTLED_SYMBOLS
settled_swap_symbol = {
    'ICP': ['2022-06-10 09:00:00', '2022-09-27 02:30:00'],
    'BNX': ['2023-02-11 04:00:00', '2023-02-22 22:45:00'],
    'TLM': ['2022-06-09 23:59:00', '2023-03-30 12:30:00']
}

settled_spot_symbol = {
    "BNX": ["2023-02-16 03:00:00", "2023-02-22 08:00:00"],
    "BTCST": ["2021-03-15 07:00:00","2021-03-19 07:00:00"],
    "COCOS": ["2021-01-19 02:00:00","2021-01-23 02:00:00"],
    "CVC": ["2022-12-09 03:00:00","2023-05-12 08:00:00"],
    "DREP": ["2021-03-29 04:00:00","2021-04-02 04:00:00"],
    "FTT": ["2022-11-15 05:00:00","2023-09-22 08:00:00"],
    "KEY": ["2023-02-10 03:00:00","2023-03-10 08:00:00"],
    "LUNA": ["2022-05-13 01:00:00","2022-05-31 06:00:00"],
    "QUICK": ["2023-07-17 03:00:00","2023-07-21 08:00:00"],
    "STRAX": ["2024-03-20 03:00:00","2024-03-28 08:00:00"],
    "SUN": ["2021-06-14 04:00:00","2021-06-18 04:00:00"],
    "VIDT": ["2022-10-31 03:00:00","2022-11-09 08:00:00"]
}

def trans_data(file_path, symbol_type):
    print(file_path)
    symbol = os.path.basename(file_path).split('-USDT.csv')[0]
    # ===跳过稳定币、杠杆代币
    if symbol.endswith(('UP', 'DOWN', 'BEAR', 'BULL')) and symbol != 'JUP' or symbol in stable_symbol:
        print(symbol, '属于不参与交易的币种，直接跳过')
        return

    benchmark = pd.DataFrame(pd.date_range(start='2017-01-01', end='2024-07-01', freq='1H'))  # 创建2017-01-01至回测结束时间的1H列表
    benchmark.rename(columns={0: 'candle_begin_time'}, inplace=True)

    df = pd.read_csv(file_path, encoding='gbk', skiprows=1, parse_dates=['candle_begin_time'])
    candle_start_time = df.iloc[0]['candle_begin_time']
    df = pd.merge(left=benchmark, right=df, on='candle_begin_time', how='left', sort=True, indicator=True)
    df = df[df['candle_begin_time'] >= candle_start_time]
    df.sort_values(by='candle_begin_time', inplace=True)
    df.drop_duplicates(subset=['candle_begin_time'], inplace=True, keep='last')
    df.reset_index(drop=True, inplace=True)
    # 数据填充
    df['close'] = df['close'].fillna(method='ffill')
    df['open'] = df['open'].fillna(df['close'])
    df['high'] = df['high'].fillna(df['close'])
    df['low'] = df['low'].fillna(df['close'])
    df['volume'] = df['volume'].fillna(0)
    df['quote_volume'] = df['quote_volume'].fillna(0)
    df['trade_num'] = df['trade_num'].fillna(0)
    df['taker_buy_base_asset_volume'] = df['taker_buy_base_asset_volume'].fillna(0)
    df['taker_buy_quote_asset_volume'] = df['taker_buy_quote_asset_volume'].fillna(0)
    df['symbol'] = df['symbol'].fillna(method='ffill').str.replace('-', '')
    df['avg_price'] = df['avg_price_1m'].fillna(df['open'])
    df['avg_price'].fillna(value=df['open'], inplace=True)
    if symbol_type == 'swap':
        df['funding_rate_raw'] = df['fundingRate']
        df['fundingRate'].fillna(method='ffill', inplace=True)
        col = ['candle_begin_time','open','high','low','close','volume','quote_volume','trade_num',
              'taker_buy_base_asset_volume','taker_buy_quote_asset_volume','avg_price','symbol',
              'fundingRate','funding_rate_raw']
        df = df[col]
        if symbol in settled_swap_symbol:
            df_old = df[df['candle_begin_time'] < settled_swap_symbol[symbol][0]].copy()
            df_old['symbol'] = f'{symbol}1USDT'
            df_old.to_feather(os.path.join(pickle_path, symbol_type, f'{symbol}1-USDT.pkl'))
            df_new = df[df['candle_begin_time'] > settled_swap_symbol[symbol][1]].copy()
            df_new.reset_index(drop=True, inplace=True)
            df_new.to_feather(os.path.join(pickle_path, symbol_type, f'{symbol}-USDT.pkl'))
            return
        df.to_feather(os.path.join(pickle_path, symbol_type, f'{symbol}-USDT.pkl'))
    if symbol_type == 'spot':
        col = ['candle_begin_time','open','high','low','close','volume','quote_volume','trade_num',
              'taker_buy_base_asset_volume','taker_buy_quote_asset_volume','avg_price','symbol']
        df = df[col]
        if symbol in settled_spot_symbol:
            df_old = df[df['candle_begin_time'] < settled_spot_symbol[symbol][0]].copy()
            df_old['symbol'] = f'{symbol}1USDT'
            df_old.to_feather(os.path.join(pickle_path, symbol_type, f'{symbol}1-USDT.pkl'))
            df_new = df[df['candle_begin_time'] > settled_spot_symbol[symbol][1]].copy()
            df_new.reset_index(drop=True, inplace=True)
            df_new.to_feather(os.path.join(pickle_path, symbol_type, f'{symbol}-USDT.pkl'))
            return
        df.to_feather(os.path.join(pickle_path, symbol_type, f'{symbol}-USDT.pkl'))

if __name__ == '__main__':
    # =====获取所有文件路径
    swap_symbol_path = glob(os.path.join(xbx_data_path , 'swap', '') + '*USDT.csv')  # 获取kline_path路径下，所有以usdt.csv结尾的文件路径
    swap_file_path = []
    for _symbol_path in swap_symbol_path:
        swap_file_path.append([_symbol_path, 'swap'])
    # 获取所有现货文件的路径
    spot_file_path = []
    if if_use_spot:
        spot_symbol_path = glob(os.path.join(xbx_data_path , 'spot', '') + '*USDT.csv')  # 获取kline_path路径下，所有以usdt.csv结尾的文件路径
        for _symbol_path in spot_symbol_path:
            spot_file_path.append([_symbol_path, 'spot'])

    # 合并合约和现货文件路径信息
    symbol_file_path = swap_file_path + spot_file_path

    # =====并行或串行，依次读取每个币种的数据，进行处理，并最终合并成一张大表输出
    multiply_process = True  # 是否并行。在测试的时候可以改成False，实际跑的时候改成True
    if multiply_process:
        Parallel(n_jobs=n_jobs)(
            delayed(trans_data)(file_path, symbol_type)
            for file_path, symbol_type in symbol_file_path)
    else:
        for file_path, symbol_type in symbol_file_path:
            trans_data(file_path, symbol_type)